
package cn.myeasyai.milvus;

import cn.myeasyai.config.MilvusConfig;
import cn.myeasyai.dto.UserDTO;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.InsertParam.Field;

import java.util.*;

import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;

/**
 * @author fushoujiang
 */
@Service
public class VectorService {
    @Resource
    private MilvusServiceClient milvusClient;

    @Resource
    private MilvusConfig milvusConfig;

    private static List<String> fieldList = Arrays.asList("face_url","face_feature_url","user_id","user_name","Auto_id");

    public List<SearchResultsWrapper.IDScore> search(List<Float> floatList,Integer topK) {
        List<List<Float>> vectors = Collections.singletonList(floatList);
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(milvusConfig.getCollectionName())
                .withMetricType(MetricType.L2)
                .withTopK(topK)
                .withOutFields(fieldList)
                .withFloatVectors(vectors)
                .withVectorFieldName(milvusConfig.getVectorFieldName()).build();
        R<SearchResults> searchResults = this.milvusClient.search(searchParam);
        if (searchResults.getStatus() == 0) {
            return  new SearchResultsWrapper(searchResults.getData().getResults()).getIDScore(0);
        }
        return Collections.emptyList();
    }


    /**
     * 初始化商家向量数据
     * @param userDTO 用户信息
     * @param features 向量数据
     */
    public void initUserVector(UserDTO userDTO, List<Float> features) {
        List<String> names = Collections.singletonList(userDTO.getUserName());
        List<Long> userIds = Collections.singletonList(userDTO.getUserId());
        List<String> getFaceUrl = Collections.singletonList(userDTO.getFaceUrl());
        List<String> getFaceFeatureUrl = Collections.singletonList(userDTO.getFaceFeatureUrl());
        List<List<Float>> vectors = Collections.singletonList(features);
        List<Field> fields = new ArrayList();
        fields.add(new Field("vector", vectors));
        fields.add(new Field("face_url", getFaceUrl));
        fields.add(new Field("face_feature_url", getFaceFeatureUrl));
        fields.add(new Field("user_id", userIds));
        fields.add(new Field("user_name", names));
        InsertParam insertParam = InsertParam.newBuilder().withCollectionName(milvusConfig.getCollectionName()).withFields(fields).build();
        this.milvusClient.insert(insertParam);
    }


    /**
     * 清空历史数据
     */
    public void cleanAll() {
        DeleteParam build = DeleteParam.newBuilder()
                .withCollectionName(milvusConfig.getCollectionName())
                .withExpr("Auto_id not in [1]")
                .build();
        this.milvusClient.delete(build);
    }

}
